package adapters

import (
	"ai-gateway/db"
	"ai-gateway/proto"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type AlibabaStreamAdapter struct{}

func (a *AlibabaStreamAdapter) StreamCall(ctx context.Context, req *proto.AIRequest,
	stream proto.AdapterService_StreamCallAIServer) error {

	// 加载阿里云配置
	var config struct {
		ApiKey     string  `db:"api_key"`
		ApiUrl     string  `db:"api_url"`
		TokenPrice float64 `db:"token_price"`
	}
	err := db.DB.Get(&config, "SELECT api_key, api_url, token_price FROM provider_configs WHERE provider = 'alibaba'")
	if err != nil {
		return status.Errorf(codes.Internal, "加载阿里云配置失败: %v", err)
	}

	// 构造标准流式请求体
	requestBody := map[string]interface{}{
		"model": "qwen-turbo", // 替换为实际模型，如 qwen-plus
		"messages": []map[string]string{
			{
				"role":    "user",
				"content": req.Prompt,
			},
		},
		"stream": true,
	}
	jsonData, _ := json.Marshal(requestBody)

	// 发起 HTTP 请求
	httpReq, err := http.NewRequestWithContext(ctx, "POST", config.ApiUrl, strings.NewReader(string(jsonData)))
	if err != nil {
		return status.Errorf(codes.Internal, "创建请求失败: %v", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+config.ApiKey)

	client := &http.Client{}
	resp, err := client.Do(httpReq)
	if err != nil {
		return status.Errorf(codes.Unavailable, "调用阿里云 API 失败: %v", err)
	}
	defer resp.Body.Close()

	// 处理流式响应（SSE格式）
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return status.Errorf(codes.Unavailable, "API 返回错误状态: %d, 响应: %s", resp.StatusCode, string(body))
	}

	reader := resp.Body
	buffer := make([]byte, 4096) // 增大缓冲区
	var accumulatedData []byte
	var tokenCount int

	for {
		n, err := reader.Read(buffer)
		if err != nil {
			if err == io.EOF {
				break
			}
			return status.Errorf(codes.Unavailable, "读取流失败: %v", err)
		}

		// 拼接响应片段
		accumulatedData = append(accumulatedData, buffer[:n]...)
		lines := strings.Split(string(accumulatedData), "\n")

		// 处理完整的 SSE 事件
		for i, line := range lines {
			if i == len(lines)-1 {
				// 剩余不完整数据，留待下次处理
				accumulatedData = []byte(line)
				break
			}
			line = strings.TrimSpace(line)
			if line == "data: [DONE]" {
				accumulatedData = nil
				break
			}
			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")
			var aiResponse struct {
				Choices []struct {
					Delta struct {
						Content string `json:"content"`
					} `json:"delta"`
					FinishReason string `json:"finish_reason"`
					Index        int    `json:"index"`
				} `json:"choices"`
			}

			if err := json.Unmarshal([]byte(data), &aiResponse); err != nil {
				continue // 跳过无效片段
			}

			// 提取流式输出
			for _, choice := range aiResponse.Choices {
				if choice.Delta.Content != "" {
					tokenCount++
					// 5. 流式推送响应
					if err := stream.Send(&proto.AIResponse{
						Result:     choice.Delta.Content,
						TokenUsed:  int32(tokenCount),
						Provider:   "alibaba",
						Cost:       int64(float64(tokenCount) * config.TokenPrice * 100),
						IsFinished: false,
					}); err != nil {
						return status.Errorf(codes.Unavailable, "推送流失败: %v", err)
					}
				}
			}
		}
	}

	// 发送结束标识
	return stream.Send(&proto.AIResponse{
		Result:     "流式响应结束", // 或实际完整结果
		TokenUsed:  int32(tokenCount),
		Cost:       int64(float64(tokenCount) * config.TokenPrice * 100),
		Provider:   "alibaba",
		IsFinished: true,
	})
}

func init() {
	RegisterStreamAdapter("alibaba", &AlibabaStreamAdapter{})
}
